package com.code.nowcoder;

import java.util.ArrayList;
import java.util.Arrays;

/**
 * @author codeDorado
 * @version 1.0
 * @date 2020/7/29 13:42
 *  dfs 并记录当前节点的子结点个数与当前节点的深度，找到初始节点到目标结点的路径即可
 */
public class The71_X {

    ArrayList<Integer>[] node;
    int depth[];
    int child[];

    /**
     * 寻找牛妹
     *
     * @param n 房间数量
     * @param m 查询数量
     * @param u u,v代表通道
     * @param v u,v代表通道
     * @param x x代表m次查询中牛妹的位置
     * @return int整型一维数组
     */
    public int[] solve(int n, int m, int[] u, int[] v, int[] x) {
        // write code here
        depth = new int[n + 1];
        child = new int[n + 1];
        node = new ArrayList[n + 1];
        int[] rtn = new int[m];
        for (int i = 1; i < n + 1; i++) {
            node[i] = new ArrayList<>();
        }
        for (int i = 0; i < n - 1; i++) {
            node[u[i]].add(v[i]);
            node[v[i]].add(u[i]);
        }

        dfs(1, -1, 0);

        int sum = (n - 1) << 1;
        for (int i = 0; i < m; i++) {
            rtn[i] = sum - child[x[i]] * 2 - depth[x[i]];
        }
        return rtn;
    }

    /**
     * @param m  索引值，表示遍历哪个 list
     * @param fa 上一个节点，如果跟遍历的节点相同
     * @param d  深度
     */
    int dfs(int m, int fa, int d) {
        depth[m] = d;
        int childs = 0;
        for (int i = 0; i < node[m].size(); i++) {
            if (node[m].get(i) != fa)
                childs += dfs(node[m].get(i), m, d + 1);
        }
        child[m] = childs;
        return childs + 1;
    }


    public static void main(String[] args) {
        The71_X the71 = new The71_X();
        System.out.println(Arrays.toString(the71.solve(5, 5, new int[]{1, 1, 2, 4}, new int[]{2, 3, 4, 5}, new int[]{5, 4, 2, 2, 4})));
    }

}
